package query

import (
	"bufio"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// MetaxQuery MetaX GPU查询实现
type MetaxQuery struct{}

func init() {
	RegisterQueryer("metax", func() GPUQuery { return &MetaxQuery{} })
}

func (q *MetaxQuery) Query() ([]GPUInfo, error) {
	// 并行获取使用率和内存信息
	var wg sync.WaitGroup
	var usageOutput, memoryOutput string
	var usageErr, memoryErr error

	wg.Add(2)
	go func() {
		defer wg.Done()
		usageOutput, usageErr = runCommand("mx-smi", "--show-usage", "-i", "all")
	}()
	go func() {
		defer wg.Done()
		memoryOutput, memoryErr = runCommand("mx-smi", "--show-memory", "-i", "all")
	}()
	wg.Wait()

	if usageErr != nil {
		return nil, fmt.Errorf("获取GPU使用率失败: %v", usageErr)
	}
	if memoryErr != nil {
		return nil, fmt.Errorf("获取GPU内存失败: %v", memoryErr)
	}

	// 解析使用率信息
	usageMap := parseMetaxUsage(usageOutput)
	// 解析内存信息
	memoryMap := parseMetaxMemory(memoryOutput)

	// 合并信息
	var gpus []GPUInfo
	for index := range usageMap {
		if memory, exists := memoryMap[index]; exists {
			gpus = append(gpus, GPUInfo{
				Index:       index,
				Utilization: usageMap[index],
				MemoryUsed:  memory.Used,
				MemoryTotal: memory.Total,
			})
		}
	}

	return gpus, nil
}

// parseMetaxUsage 解析MetaX GPU使用率输出
func parseMetaxUsage(output string) map[int]int {
	result := make(map[int]int)
	scanner := bufio.NewScanner(strings.NewReader(output))

	// 正则表达式匹配GPU索引和使用率
	gpuRegex := regexp.MustCompile(`^GPU#(\d+)`)
	utilRegex := regexp.MustCompile(`GPU\s+:\s+(\d+) %`)

	currentGPU := -1
	for scanner.Scan() {
		line := scanner.Text()

		// 检查是否是GPU行
		if matches := gpuRegex.FindStringSubmatch(line); matches != nil {
			if index, err := strconv.Atoi(matches[1]); err == nil {
				currentGPU = index
			}
			continue
		}

		// 检查是否是使用率行
		if currentGPU != -1 {
			if matches := utilRegex.FindStringSubmatch(line); matches != nil {
				if utilization, err := strconv.Atoi(matches[1]); err == nil {
					result[currentGPU] = utilization
					currentGPU = -1 // 重置当前GPU
				}
			}
		}
	}

	return result
}


// parseMetaxMemory 解析MetaX GPU内存输出
func parseMetaxMemory(output string) map[int]MemoryInfo {
	result := make(map[int]MemoryInfo)
	scanner := bufio.NewScanner(strings.NewReader(output))

	// 正则表达式匹配GPU索引和内存信息
	gpuRegex := regexp.MustCompile(`^GPU#(\d+)`)
	memoryRegex := regexp.MustCompile(`vram (used|total)\s+:\s+(\d+) KB`)

	currentGPU := -1
	memory := MemoryInfo{}

	for scanner.Scan() {
		line := scanner.Text()

		// 检查是否是GPU行
		if matches := gpuRegex.FindStringSubmatch(line); matches != nil {
			// 保存前一个GPU的信息
			if currentGPU != -1 {
				result[currentGPU] = memory
				memory = MemoryInfo{}
			}

			if index, err := strconv.Atoi(matches[1]); err == nil {
				currentGPU = index
			}
			continue
		}

		// 检查是否是内存信息行
		if currentGPU != -1 {
			if matches := memoryRegex.FindStringSubmatch(line); matches != nil {
				value, err := strconv.Atoi(matches[2])
				if err != nil {
					continue
				}

				// 转换为MB (除以1024)
				valueMB := value / 1024

				if matches[1] == "used" {
					memory.Used = valueMB
				} else if matches[1] == "total" {
					memory.Total = valueMB
				}
			}
		}
	}

	// 保存最后一个GPU的信息
	if currentGPU != -1 {
		result[currentGPU] = memory
	}

	return result
}
